﻿// convert CMA forcast Wind data from grib2 to GeoTiff. 2025/5/9
// example filename: Z_NWGD_C_BABJ_20250505150023_P_RFFC_SCMOC-WWD01_202505052000_07201.GRB2
// it is a forcast data for every hour in next 72 hours.

#include <iostream>
#include "gdal_priv.h"
#include "grib2.h"
#include <string>
#include <vector>
#include <memory>
using namespace std;
 

int main(int argc,char* argv[])
{
    cout << "A program to convert CMA forcast Wind data from grib2 to GeoTiff." << endl;
    cout << "The result band order is u1,v1,u2,v2...."<<endl;
    cout<<"v1.0.0"<<endl;//2025/5/9
    cout<<"usage:yy_grib2winddecode /some/Z_...GRB2 /out/dir/filename.tif"<<endl;

    if( argc!=3 ) {
        cout<<"argc not 3"<<endl;
        return 11 ;
    }

    GDALAllRegister();

    string infilename = argv[1] ;
    // "/home/hadoop/tempdata/Z_NWGD_C_BABJ_20250505150023_P_RFFC_SCMOC-WWD01_202505052000_07201.GRB2" ;
    string outfilename = argv[2];//"test-out.tif" ;

    const int WID = 1401 ;
    const int HEI = 1201 ;
    const int ASIZE = WID*HEI;
    
    unsigned char *cgrib;
    g2int  listsec0[3],listsec1[13],numlocal,numfields;
    long   lskip,n,lgrib,iseek,expand;
    int    unpack,ret,ierr;
    gribfield  *gfld;
    FILE   *fptr;
    size_t  lengrib;
 
    iseek=0;
    unpack=1;
    expand=1;
    int nmessage=0;
    fptr=fopen( infilename.c_str() ,"r");
    if(true) { // for (;;) {
         seekgb(fptr,iseek,32000,&lskip,&lgrib);
         if (lgrib == 0) {
            cout<<"Error, end loop at EOF or problem."<<endl;
            return 22 ;
         } // end loop at EOF or problem
         cgrib=(unsigned char *)malloc(lgrib);
         ret=fseek(fptr,lskip,SEEK_SET);
         lengrib=fread(cgrib,sizeof(unsigned char),lgrib,fptr);
         iseek=lskip+lgrib;
         ierr=g2_info(cgrib,listsec0,listsec1,&numfields,&numlocal);
         cout<<"numfields "<<numfields<<endl;

        CPLStringList opt ;
        opt.SetNameValue("COMPRESS" , "Deflate") ;

        GDALDriver *driver = GetGDALDriverManager()->GetDriverByName( "GTiff" );
        GDALDataset* dsout = driver->Create( outfilename.c_str() , 
            WID , HEI , numfields ,GDT_Float32,  opt.List()) ;

        for (n=0;n<numfields;n++) {
            ierr=g2_getfld(cgrib,n+1,unpack,expand,&gfld);
            cout<<"ifield " << n <<endl;
            // cout<<"ndpts "<<gfld->ndpts<<endl;//1401(cols)x1201(rows) = 1682601
            if( gfld->ndpts == WID*HEI ) {
                vector<float> buffer(WID*HEI);
                for(int iy=0;iy<HEI;++iy) {
                    for(int ix = 0 ; ix < WID ; ++ ix ) {
                        //invertY
                        int it = iy*WID+ix ;
                        int it2 = (HEI-1-iy)*WID+ix ;
                        buffer[it2] = gfld->fld[it] ;
                    }
                }
                dsout->GetRasterBand(n+1)->RasterIO(
                    GF_Write,0,0,WID,HEI,buffer.data(),WID,HEI,GDT_Float32,0,0,0) ;

            }else {
                cout<<"bad ndpts not equal with "<<WID*HEI<<endl;
            }

            g2_free(gfld);
        }
        free(cgrib);
        // I fix all geotrans and proj info here for convient and fast use. I know its stupid.
        double trans[6] = {69.975, 0.05, 0.0, 60.025, 0.0, -0.05 } ;
        dsout->SetGeoTransform(trans);
        dsout->SetProjection("GEOGCS[\"WGS 84\",DATUM[\"WGS_1984\",SPHEROID[\"WGS 84\","
            "6378137,298.257223563,AUTHORITY[\"EPSG\",\"7030\"]],AUTHORITY[\"EPSG\",\"6326\"]]"
            ",PRIMEM[\"Greenwich\",0,AUTHORITY[\"EPSG\",\"8901\"]],UNIT[\"degree\","
            "0.01745329251994328,AUTHORITY[\"EPSG\",\"9122\"]],AUTHORITY[\"EPSG\",\"4326\"]]");  
        GDALClose(dsout);
        cout<<"write to "<<outfilename<<endl;
    }

    cout<<"done"<<endl;
    return 0;
}
